"""
音频数据增强增强版
功能：带参数校验、多进程处理、自动跳过已增强文件、完善日志系统的批量音频增强工具
"""

import librosa
import numpy as np
from pathlib import Path
import yaml
import soundfile as sf
from audiomentations import Compose, AddGaussianNoise, PitchShift, TimeStretch
import multiprocessing
import logging
from tqdm import tqdm
from typing import List, Dict

# 配置日志系统
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        logging.FileHandler("logs/audio_augmentation.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


def load_config(config_path: Path) -> Dict:
    """
    加载并验证配置文件

    参数:
        config_path: 配置文件路径

    返回:
        配置字典

    异常:
        ValueError: 当缺少必要参数时
    """
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            cfg = yaml.safe_load(f)

        # 验证必要参数
        required_keys = {"sample_rate", "augment_copies", "data_dir"}
        if not required_keys.issubset(cfg.keys()):
            missing = required_keys - cfg.keys()
            raise ValueError(f"缺少必要配置参数: {missing}")

        return cfg
    except Exception as e:
        logger.error(f"配置文件加载失败: {str(e)}")
        raise


def create_augmentation_pipeline(cfg: Dict) -> Compose:
    """
    创建音频增强流水线

    参数:
        cfg: 配置字典

    返回:
        配置好的增强流水线对象
    """
    return Compose([
        # 高斯噪声增强 (50%概率应用)
        AddGaussianNoise(
            min_amplitude=cfg.get("noise_min_amp", 0.001),
            max_amplitude=cfg.get("noise_max_amp", 0.015),
            p=0.5
        ),
        # 音高偏移 (-2到2个半音, 50%概率)
        PitchShift(
            min_semitones=cfg.get("pitch_min", -2),
            max_semitones=cfg.get("pitch_max", 2),
            p=0.5
        ),
        # 时间拉伸 (0.8x-1.2x速度, 50%概率)
        TimeStretch(
            min_rate=cfg.get("time_min_rate", 0.8),
            max_rate=cfg.get("time_max_rate", 1.2),
            p=0.5
        )
    ])


def process_single_file(args: tuple) -> None:
    """
    处理单个音频文件（多进程 worker 函数）

    参数:
        args: 包含以下参数的元组:
            - wav_path: 音频文件路径
            - output_dir: 输出目录
            - augment: 增强流水线
            - cfg: 配置字典
            - file_counter: 进程安全计数器
    """
    wav_path, output_dir, augment, cfg = args
    try:
        # 检查是否已存在增强文件
        existing_files = list(output_dir.glob(f"{wav_path.stem}_aug*.wav"))
        if len(existing_files) >= cfg["augment_copies"]:
            return

        # 加载音频 (使用soundfile保持一致性)
        y, sr = sf.read(wav_path, dtype="float32")
        if y.ndim > 1:
            y = librosa.to_mono(y.T)

        # 重采样检查
        if sr != cfg["sample_rate"]:
            y = librosa.resample(y, orig_sr=sr, target_sr=cfg["sample_rate"])
            sr = cfg["sample_rate"]

        # 生成增强样本
        for i in range(cfg["augment_copies"]):
            output_path = output_dir / f"{wav_path.stem}_aug{i}.wav"
            if not output_path.exists():
                augmented = augment(samples=y, sample_rate=sr)
                sf.write(output_path, augmented, sr)

    except Exception as e:
        logger.error(f"处理文件失败 {wav_path}: {str(e)}", exc_info=True)


def augment_dataset(cfg: Dict) -> None:
    """
    主增强函数（最终修复版）
    """
    data_dir = Path(cfg["data_dir"])
    raw_dir = data_dir / "raw"
    augmented_dir = data_dir / "augmented"

    # 创建输出目录
    augmented_dir.mkdir(parents=True, exist_ok=True)

    # 初始化增强流水线
    augment = create_augmentation_pipeline(cfg)

    # 准备任务列表
    tasks = []
    for class_dir in raw_dir.iterdir():
        if class_dir.is_dir():
            class_output_dir = augmented_dir / class_dir.name
            class_output_dir.mkdir(exist_ok=True)

            for wav in class_dir.glob("*.wav"):
                tasks.append((wav, class_output_dir, augment, cfg))

    # 使用Manager正确创建共享计数器
    manager = multiprocessing.Manager()
    counter = manager.Value('i', 0)
    lock = manager.Lock()

    # 定义回调函数
    def update_counter(_):
        nonlocal counter, lock
        with lock:
            counter.value += 1

    # 多进程处理
    with multiprocessing.Pool(processes=cfg.get("num_workers", 4)) as pool:
        results = []

        # 提交任务
        for task in tasks:
            res = pool.apply_async(
                process_single_file,
                args=(task,),
                callback=update_counter
            )
            results.append(res)

        # 进度条监控
        with tqdm(total=len(tasks), desc="增强进度") as pbar:
            for res in results:
                res.wait()
                pbar.update(1)

    logger.info(f"完成! 共处理 {counter.value} 个文件，生成 {counter.value * cfg['augment_copies']} 个增强样本")
    manager.shutdown()  # 清理Manager资源

if __name__ == "__main__":
    try:
        # 加载配置
        cfg = load_config(Path("configs/default.yaml"))

        # 记录配置
        logger.info("启动音频增强任务")
        logger.info(f"配置参数:\n{yaml.dump(cfg, indent=2)}")

        # 执行增强
        augment_dataset(cfg)
        logger.info("任务成功完成")

    except Exception as e:
        logger.error(f"任务异常终止: {str(e)}", exc_info=True)
        exit(1)